import json
import random
from PIL import Image
from environment.stretch_controller import StretchController
import prior
from matplotlib import pyplot as plt

from typing import List

def draw(expert, name="./debug/init.png"):
        fig, ax = plt.subplots(1, 1)
                
        xs_obs = [obs["x"] for obs in expert.obsmap_positions]
        zs_obs = [obs["z"] for obs in expert.obsmap_positions]
        plt.plot(xs_obs, zs_obs, 'sk')
        plt.plot(expert.pathx, expert.pathz, '-r')
        plt.plot(expert.start_node.x, expert.start_node.z, 'sg')
        plt.plot(expert.goal_node.x, expert.goal_node.z, 'sb')
        
        plt.axis("equal")
        plt.show()
        plt.savefig(f"{name}.png")

def l2_distance(p1, p2):
    """Calculate the L2 distance between two points."""
    return ((p1['x'] - p2['x']) ** 2 + (p1['z'] - p2['z']) ** 2) ** 0.5


def load_procthor_houses():
    """Load the Procthor houses dataset."""
    return prior.load_dataset(dataset="spoc-data", entity="spoc-robot", revision="houses-test-val")["val"]

def get_position_rotation(controller:StretchController):
    """Get the current position and rotation of the robot."""
    pos_ro = controller.get_current_agent_full_pose()
    return [pos_ro['position']['x'],pos_ro['position']['y'],pos_ro['position']['z'],
                              pos_ro['rotation']['x'],pos_ro['rotation']['y'],pos_ro['rotation']['z']]


def filter_objects(controller:StretchController, object_types:List):
    """Filter objects based on target types."""
    # random.seed(0)
    objects = [o for o in controller.get_objects() if o["objectType"] in object_types]
    random.shuffle(objects)
    return objects


def target_types(datapath="./expert/domain/procthor.yaml"):
    """Load target object types from a YAML file."""
    import yaml
    with open(datapath, 'r') as file:
        data = yaml.safe_load(file)
    return data.get('target_object_types', [])

def get_one_object(controller:StretchController,id="Bowl|+00.79|+00.90|-00.12"):
    object_types = target_types()
    objects = filter_objects(controller, object_types)
    print(f"Found {len(objects)} objects of type {object_types}")
    if id:
        objects = [o for o in objects if o["objectId"] == id]
        return None if len(objects) == 0 else objects
    return None if len(objects) == 0 else random.choice(objects)

def save_navigation_frame(nav_frames, frame_dir):
    """Save navigation frames to a directory."""
    for i, nav_frame in enumerate(nav_frames):
        nav_frame = Image.fromarray(nav_frame)
        if nav_frame.mode != "RGB":
            nav_frame = nav_frame.convert("RGB")
        nav_frame.save(f"{frame_dir}/nav_frame_{i}.png")


def save_data_json(data, save_path, should_save=True):
    """Save data to a file if should_save is True."""
    if should_save:
        try:
            with open(save_path, "w") as f:
                json.dump(data, f, indent=4)
        except IOError as e:
            print(f"Error saving data to {save_path}: {e}")

def positions2path(positions):
    """Convert positions to a path format."""
    return [{'x': point[0], 'y': point[1], 'z': point[2]} for point in positions]

def get_top_down_map(controller:StretchController,positions, save_path=None):
    """Get a top-down map of the environment."""
    shortestpath = positions2path(positions)
    top_down_map = controller.get_top_down_path_view(shortestpath)
    top_down_map = Image.fromarray(top_down_map)
    if top_down_map.mode != "RGB":
        top_down_map = top_down_map.convert("RGB")
    if save_path:
        top_down_map.save(save_path)
    return top_down_map